package io.growing.adaptor.common.async;

import java.util.*;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.ReentrantLock;
import java.util.stream.Collectors;

/**
 * @ClassName PriorityDistributedBlockingQueue @Description 根据不同的优先级分发到不同的阻塞队列，数字越小，级别越高，越最先被消费执行
 * 作用是将Runnble按优先级区分，比如在查询用户属性的业务上，所有查询用户属性的操作都共用一个线程池，但保证接口调用优先于kafka消费端执行，
 * 在kafka消费端又分为用户属性查询和埋点上报，埋点是最低优先级 @Author wangjiu @Date 2021/12/8 2:32 下午
 */
public class PriorityDistributedBlockingQueue extends AbstractQueue<Runnable>
        implements BlockingQueue<Runnable>, java.io.Serializable {

    private final Map<Integer, BlockingQueue<Runnable>> queueMap;

    /** 第一队列是最主要的队列，线程池最主要就是从这个队列取任务，无论什么优先级的任务都可能被放进第一队列，map里的其他队列其实不是很重要 */
    private final BlockingQueue<Runnable> firstQueue;

    /** lastQueue是为了兜底， */
    private final BlockingQueue<Runnable> lastQueue;

    private final List<Map.Entry<Integer, BlockingQueue<Runnable>>> queueEntryList;

    private final List<BlockingQueue<Runnable>> queueList;

    /** 表示在第一个队列里的最高优先级是多少，只要<=这个优先级，都可以往第一队列里放 */
    private final AtomicInteger minPriority = new AtomicInteger(Integer.MAX_VALUE);

    private final ReentrantLock lock = new ReentrantLock(true);

    /**
     * 目前的业务需求(查询用户属性)，只有一个队列也能做到优先级判断，这里用一个map,是为了能扩展
     *
     * @param queueMap
     */
    public PriorityDistributedBlockingQueue(Map<Integer, BlockingQueue<Runnable>> queueMap) {

        this.queueMap = queueMap;
        queueEntryList =
                queueMap.entrySet().stream()
                        .sorted(Comparator.comparingInt(Map.Entry::getKey))
                        .collect(Collectors.toList());
        queueList = queueEntryList.stream().map(Map.Entry::getValue).collect(Collectors.toList());
        firstQueue = queueList.get(0);
        lastQueue = queueList.get(queueList.size() - 1);
    }

    /**
     * 默认全放到第一个队列里，就是最高优先级
     *
     * @param runnable
     * @return
     */
    @Override
    public boolean offer(Runnable runnable) {
        lock.lock();
        try {
            if (runnable instanceof PriorityRunnable) {
                int priority = ((PriorityRunnable) runnable).getPriority();
                // 更新优先级
                int curMinPriority = updateMinPriority(priority);
                if (priority <= curMinPriority) {
                    // 虽然上面一行更新了优先级，但是这一行依然可能会失败，比如被其他线程抢先放了任务进去，
                    // 但是无所谓，更新优先级才是最重要的目的，这里能放进队列就放，放不进去就算了，
                    // 因为优先级更新以后，以后更低优先级的任务就放不进去了，后续同优先级或高优先级的任务依然有机会放进去，而当前任务被拒绝以后可以自己去执行
                    return firstQueue.offer(runnable);
                } else {
                    return offer(runnable, priority);
                }
            } else {
                // 没有设置优先级的任务按最低优先级执行，如果有高优先级的任务，那就没机会放到队列里了
                return getMinPriority() == Integer.MAX_VALUE && firstQueue.offer(runnable);
            }
        } finally {
            lock.unlock();
        }
    }

    @Override
    public boolean offer(Runnable runnable, long timeout, TimeUnit unit)
            throws InterruptedException {
        lock.lockInterruptibly();
        try {
            if (runnable instanceof PriorityRunnable) {
                int priority = ((PriorityRunnable) runnable).getPriority();
                // 更新优先级
                int curMinPriority = updateMinPriority(priority);
                if (priority <= curMinPriority) {
                    // 虽然上面一行更新了优先级，当时这一行依然可能会失败，比如被其他线程抢先放了任务进去，
                    // 但是无所谓，更新优先级才是最重要的目的，这里能放进队列就放，放不进去就算了，
                    // 因为优先级更新以后，以后更低优先级的任务就放不进去了，后续同优先级或高优先级的任务就有机会放进去了，而当前任务被拒绝以后可以自己去执行
                    return firstQueue.offer(runnable, timeout, unit);
                } else {
                    return offer(runnable, priority);
                }
            } else {
                return getMinPriority() == Integer.MAX_VALUE
                        && firstQueue.offer(runnable, timeout, unit);
            }
        } finally {
            lock.unlock();
        }
    }

    /**
     * 将任务按照不同的优先级放到不同的队列中，如果放不进去就返回false,
     *
     * @param runnable
     * @param priority
     * @return
     */
    public boolean offer(Runnable runnable, int priority) {
        boolean r = false;
        lock.lock();
        try {
            for (Map.Entry<Integer, BlockingQueue<Runnable>> e : queueEntryList) {
                if (priority <= e.getKey()) {
                    BlockingQueue queue = e.getValue();
                    if (r = queue.offer(runnable)) {
                        if (queue.equals(firstQueue)) {
                            // 如果刚好放到了第一个队列里，就尝试更新一下最高优先级
                            updateMinPriority(priority);
                        }
                        return r;
                    }
                }
            }
            return r || lastQueue.offer(runnable);
        } finally {
            lock.unlock();
        }
    }

    public int getMinPriority() {
        return minPriority.get();
    }

    @Override
    public Iterator<Runnable> iterator() {
        return firstQueue.iterator();
    }

    @Override
    public int size() {
        return queueList.stream().mapToInt(BlockingQueue::size).sum();
    }

    @Override
    public void put(Runnable runnable) throws InterruptedException {
        firstQueue.put(runnable);
    }

    public void put(Runnable runnable, int priority) throws InterruptedException {
        for (Map.Entry<Integer, BlockingQueue<Runnable>> e : queueEntryList) {
            if (priority <= e.getKey()) {
                e.getValue().put(runnable);
                return;
            }
        }
        lastQueue.put(runnable);
    }

    /**
     * 不断的遍历每个队列，直到能从队列里取出一个任务，遍历队列时的阻塞时间可以改成可配的
     *
     * @return
     * @throws InterruptedException
     */
    @Override
    public Runnable take() throws InterruptedException {
        Runnable r = null;
        while (r == null) {
            for (BlockingQueue<Runnable> q : queueList) {
                if ((r = q.poll(200, TimeUnit.MILLISECONDS)) != null) {
                    break;
                }
            }
        }
        resetMinPriority();
        return r;
    }

    @Override
    public Runnable poll(long timeout, TimeUnit unit) throws InterruptedException {
        Runnable r = null;
        for (BlockingQueue<Runnable> q : queueList) {
            if ((r = q.poll(timeout, unit)) != null) {
                break;
            }
        }
        resetMinPriority();
        return r;
    }

    @Override
    public int remainingCapacity() {
        long sum = queueList.stream().mapToLong(BlockingQueue::remainingCapacity).sum();
        if (sum > Integer.MAX_VALUE) {
            return Integer.MAX_VALUE;
        }
        return (int) sum;
    }

    @Override
    public int drainTo(Collection<? super Runnable> c) {
        int sum = queueList.stream().mapToInt(q -> q.drainTo(c)).sum();
        resetMinPriority();
        return sum;
    }

    @Override
    public int drainTo(Collection<? super Runnable> c, int maxElements) {
        int sum = queueList.stream().mapToInt(q -> q.drainTo(c, maxElements)).sum();
        resetMinPriority();
        return sum;
    }

    @Override
    public Runnable poll() {
        Runnable r = null;
        for (BlockingQueue<Runnable> q : queueList) {
            if ((r = q.poll()) != null) {
                break;
            }
        }
        resetMinPriority();
        return r;
    }

    @Override
    public Runnable peek() {
        Runnable r = null;
        for (BlockingQueue<Runnable> q : queueList) {
            if ((r = q.peek()) != null) {
                break;
            }
        }
        return r;
    }

    /**
     * 更新优先级并返回更新后的值
     *
     * @param priority
     * @return
     */
    private int updateMinPriority(int priority) {
        int oldMinPriority = getMinPriority();
        while (oldMinPriority > priority && !minPriority.compareAndSet(oldMinPriority, priority)) {
            oldMinPriority = getMinPriority();
        }
        return getMinPriority();
    }

    private void resetMinPriority() {
        // 因为不确定第一队列里的任务里面的最高优先级是多少，也没必要去确定，
        // 因为要确定的话，这里只能去加锁，然后遍历第一队列，否则不加锁的话不能保证自己遍历的时候没有后插如的任务，就不能保证自己得到的优先级是真实的优先级，
        // 这里等第一队列清空后直接重置，所以第一队列不应该设置的太大，太大的话这里会等较长时间才能被重置，影响后续低优先级任务的执行
        if (firstQueue.isEmpty() && getMinPriority() < Integer.MAX_VALUE) {
            lock.lock();
            try {
                if (firstQueue.isEmpty()) {
                    minPriority.set(Integer.MAX_VALUE);
                }
            } finally {
                lock.unlock();
            }
        }
    }
}
